"""
集成增强记忆功能的智能体类
完整集成所有增强记忆组件到现有AgentClass中
"""

import os
from typing import Dict, Any, Optional, List
import logging

from langgraph.checkpoint.memory import MemorySaver

from .enhanced_memory_class import EnhancedMemoryClass

logger = logging.getLogger(__name__)

class EnhancedAgentClass:
    """增强的智能体类（简化版）"""

    def __init__(self, config: Dict[str, Any] = None):
        """初始化增强智能体类（简化版）"""
        # 基础配置
        self.modelname = os.getenv("BASE_MODEL", "gpt-4")
        self.memorykey = os.getenv("MEMORY_KEY", "chat_history")
        self.config = config or {}
        
        # 初始化ChatOpenAI模型实例
        from langchain_openai import ChatOpenAI
        self.model = ChatOpenAI(
            model=self.modelname,
            api_key=os.getenv("OPENAI_API_KEY"),
            base_url=os.getenv("OPENAI_BASE_URL")
        )
        
        # 初始化增强记忆系统
        self.enhanced_memory = EnhancedMemoryClass(
            memoryKey=self.memorykey, 
            model=self.modelname,
            config=self.config
        )
        
        # 初始化性能优化组件（简化版）
        self.cache_manager = None

        # LangGraph组件（简化版）
        self.checkpointer = MemorySaver()
        
        # 工作流相关
        self.workflow = None
        self.app = None
        
        # 初始化标志
        self._initialized = False
        
        logger.info("Enhanced Agent Class initialized")
    
    def initialize(self):
        """同步初始化所有组件（简化版）"""
        if self._initialized:
            return

        try:
            # 1. 确保增强记忆已初始化（同步方式）
            if hasattr(self.enhanced_memory, 'ensure_initialized'):
                self.enhanced_memory.ensure_initialized()

            # 2. 缓存管理器已简化，不再使用复杂的HybridCacheManager
            self.cache_manager = None  # 简化版本不使用复杂缓存

            # 3. 初始化Text2SQL智能体
            try:
                from text2sql.text2sql_tools import get_text2sql_agent
                self.text2sql_agent = get_text2sql_agent().get_agent()
            except Exception as e:
                logger.warning(f"Failed to initialize text2sql agent: {e}")
                self.text2sql_agent = None

            # 4. 使用基础工作流，避免复杂的异步初始化
            self.workflow = None  # 简化版本不使用复杂工作流

            self._initialized = True
            logger.info("Enhanced Agent Class initialized (simplified mode)")

        except Exception as e:
            logger.error(f"Failed to initialize enhanced agent: {e}")
            # 降级到基础模式
            self.workflow = None
    

    
    def run_agent(self, input_text: str, user_id: str = None) -> Dict[str, Any]:
        """运行增强智能体（简化版）"""
        try:
            # 确保已初始化
            if not self._initialized:
                self.initialize()

            # 获取用户ID
            if user_id is None:
                from dingding.Storage import get_user
                user_id = get_user("userid") or "default_user"

            session_id = f"user_{user_id}"

            # 简化的结果处理，避免复杂的异步工作流
            try:
                # 如果有工作流，尝试使用
                if hasattr(self, 'workflow') and self.workflow:
                    # 这里需要同步版本的工作流调用
                    # 暂时返回基础响应
                    output = f"收到您的消息：{input_text}。系统正在简化模式下运行。"
                else:
                    output = f"收到您的消息：{input_text}。系统正在基础模式下运行。"
            except Exception as workflow_e:
                logger.warning(f"Workflow execution failed: {workflow_e}")
                output = f"收到您的消息：{input_text}。"

            # 存储对话到增强记忆（使用传统方式）
            try:
                self.enhanced_memory.store_conversation_turn(
                    input_text, output, user_id, session_id
                )
            except Exception as store_e:
                logger.warning(f"Failed to store conversation: {store_e}")

            # 同时保存到传统记忆系统（向后兼容）
            try:
                if hasattr(self, 'basic_agent') and self.basic_agent:
                    self.basic_agent._save_to_memory(session_id, input_text, output)
                    logger.info("已同步保存到传统记忆系统")
            except Exception as e:
                logger.warning(f"保存到传统记忆系统失败: {e}")

            return {
                "output": output,
                "user_id": user_id,
                "session_id": session_id,
                "success": True,
                "memory_info": {
                    "rewrite_applied": False,
                    "context_used": 0,
                    "profile_available": False
                }
            }
            
        except Exception as e:
            logger.error(f"Error in enhanced run_agent: {e}")
            return {
                "output": f"抱歉，处理您的请求时出现了错误：{str(e)}",
                "user_id": user_id or "unknown",
                "success": False,
                "error": str(e)
            }
    
    def _extract_ai_response(self, result: Dict[str, Any]) -> str:
        """提取AI回复（简化版）"""
        try:
            if isinstance(result, dict):
                if "output" in result:
                    return result["output"]
                elif "messages" in result and result["messages"]:
                    # 简化的消息提取
                    last_message = result["messages"][-1]
                    if hasattr(last_message, 'content'):
                        return last_message.content
                    else:
                        return str(last_message)

            return "抱歉，我无法处理您的请求。"

        except Exception as e:
            logger.error(f"Failed to extract AI response: {e}")
            return "抱歉，处理响应时出现错误。"
    
    def get_conversation_history(self, user_id: str = None) -> List[Dict[str, str]]:
        """获取对话历史（同步版本）"""
        try:
            if user_id is None:
                from dingding.Storage import get_user
                user_id = get_user("userid") or "default_user"

            session_id = f"user_{user_id}"

            # 使用传统方式获取历史
            return self.enhanced_memory._get_traditional_history(session_id)

        except Exception as e:
            logger.error(f"Error getting conversation history: {e}")
            return []

    def clear_memory(self, user_id: str = None) -> bool:
        """清除用户记忆（同步版本）"""
        try:
            if user_id is None:
                from dingding.Storage import get_user
                user_id = get_user("userid") or "default_user"

            session_id = f"user_{user_id}"

            # 清除传统记忆
            if hasattr(self.enhanced_memory, 'memory_dict'):
                if session_id in self.enhanced_memory.memory_dict:
                    del self.enhanced_memory.memory_dict[session_id]

            logger.info(f"Memory cleared for user: {user_id}")
            return True

        except Exception as e:
            logger.error(f"Error clearing memory: {e}")
            return False
    
    def get_performance_stats(self) -> Dict[str, Any]:
        """获取性能统计（同步版本）"""
        try:
            stats = {
                "initialized": self._initialized,
                "memory_system": "simplified" if self._initialized else "basic"
            }

            # 添加记忆统计
            if self._initialized and self.enhanced_memory:
                try:
                    # 获取传统记忆统计
                    memory_dict = getattr(self.enhanced_memory, 'memory_dict', {})
                    stats["memory_sessions"] = len(memory_dict)
                    stats["total_conversations"] = sum(
                        len(history) if hasattr(history, '__len__') else 0
                        for history in memory_dict.values()
                    )
                except Exception as e:
                    logger.warning(f"Failed to get memory stats: {e}")
                    stats["memory_sessions"] = 0
                    stats["total_conversations"] = 0
            else:
                stats["memory_sessions"] = 0
                stats["total_conversations"] = 0

            return stats

        except Exception as e:
            logger.error(f"Error getting performance stats: {e}")
            return {"error": str(e)}

# 工厂函数
def create_enhanced_agent(config: Dict[str, Any] = None) -> EnhancedAgentClass:
    """创建增强智能体实例（简化版）"""
    agent = EnhancedAgentClass(config)
    agent.initialize()  # 同步初始化
    return agent
